
'''
auto pipe
'''
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler


model_id = "stabilityai/stable-diffusion-2-1"
model_id = "D:\code\other\stablediffusion\huggingface\hub\models--CompVis--stable-diffusion-v1-4"
# model_id = r"D:\code\other\stablediffusion\v2-1_768-ema-pruned.ckpt"
# Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.scheduler.set_timesteps(2)
print(f'timesteps {pipe.scheduler.timesteps}')
# pipe.enable_attention_slicing()
# pipe = pipe.to("cuda")
# pipe = pipe.to("cpu")
# pipe.enable_attention_slicing()

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt,num_inference_steps=2).images[0]

image.save("astronaut_rides_horse2.png")



# from diffusers import DDPMScheduler, UNet2DModel
#
# model_id = "stabilityai/stable-diffusion-2-1"
# scheduler = DDPMScheduler.from_pretrained(model_id)
# model = UNet2DModel.from_pretrained(model_id, use_safetensors=True).to("cuda")

